import torchvision
import torch.nn as nn
import torch
import timm
class Mymodel(nn.Module):
    def __init__(self,pretain=True,num_class=6000,model_name='swin_base_patch4_window7_224'):
        super(Mymodel, self).__init__()
        self.model = timm.create_model(model_name,pretrained=pretain,num_classes=num_class)

    def forward(self,x):
        out = self.model(x)
        return out


def get_img_output_length(width, height):
    def get_output_length(input_length):
        # input_length += 6
        filter_sizes = [2, 2, 2, 2, 2]
        padding = [0, 0, 0, 0, 0]
        stride = 2
        for i in range(5):
            input_length = (input_length + 2 * padding[i] - filter_sizes[i]) // stride + 1
        return input_length

    return get_output_length(width) * get_output_length(height)


class Siamese(nn.Module):
    def __init__(self, pretrained=False):
        super(Siamese, self).__init__()
        self.model = timm.create_model(model_name='swin_base_patch4_window7_224', pretrained=False)


        #flat_shape = 512 * get_img_output_length(input_shape[1], input_shape[0])
        self.fully_connect1 = torch.nn.Linear(1024, 512)
        self.fully_connect2 = torch.nn.Linear(512, 1)

    def forward(self, x1,x2):

        x1 = self.model.forward_features(x1)
        x2 = self.model.forward_features(x2)
        x1= torch.flatten(x1,1)
        x2 = torch.flatten(x2, 1)
        x = torch.abs(x1 - x2)
        x = self.fully_connect1(x)
        x = self.fully_connect2(x)
        x = nn.Sigmoid()(x)
        return x

if __name__ == '__main__':
    x= torch.randn(4,3,224,224)
    y= torch.randn(4,3,224,224)
    #flat_shape = 512 * get_img_output_length(224, 224)
    #print(flat_shape)
    model = Siamese()
    #model = Mymodel(model_name='resnet50', pretain=False)
    #model = del global_pool
    #model = timm.create_model(model_name='resnet50', pretrained=False,num_classes=0, global_pool='')
    print(model(x,y))
    #print(model.forward_features(x).shape)
    #model = Mymodel(num_class=10,model_name='convit_base',pretain=False)
    #print(model(x))
    #print(timm.list_models(pretrained=True))